package io.debezium.connector.postgresql;

import io.debezium.connector.SnapshotRecord;
import io.debezium.connector.postgresql.connection.GaussDBConnection;
import io.debezium.connector.postgresql.connection.Lsn;
import io.debezium.connector.postgresql.spi.OffsetState;
import io.debezium.pipeline.source.snapshot.incremental.IncrementalSnapshotContext;
import io.debezium.pipeline.source.snapshot.incremental.SignalBasedIncrementalSnapshotContext;
import io.debezium.pipeline.spi.OffsetContext;
import io.debezium.pipeline.txmetadata.TransactionContext;
import io.debezium.relational.TableId;
import io.debezium.schema.DataCollectionId;
import io.debezium.time.Conversions;
import io.debezium.util.Clock;
import org.apache.kafka.connect.data.Schema;
import org.apache.kafka.connect.data.Struct;
import org.apache.kafka.connect.errors.ConnectException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.sql.SQLException;
import java.time.Instant;
import java.util.HashMap;
import java.util.Map;

public class PostgresOffsetContext implements OffsetContext {
    private static final Logger LOGGER =
            LoggerFactory.getLogger(PostgresSnapshotChangeEventSource.class);
    public static final String LAST_COMPLETELY_PROCESSED_LSN_KEY = "lsn_proc";
    public static final String LAST_COMMIT_LSN_KEY = "lsn_commit";
    private final Schema sourceInfoSchema;
    private final SourceInfo sourceInfo;
    private boolean lastSnapshotRecord;
    private Lsn lastCompletelyProcessedLsn;
    private Lsn lastCommitLsn;
    private Lsn streamingStoppingLsn;
    private final TransactionContext transactionContext;
    private final IncrementalSnapshotContext<TableId> incrementalSnapshotContext;

    private PostgresOffsetContext(
            PostgresConnectorConfig connectorConfig,
            Lsn lsn,
            Lsn lastCompletelyProcessedLsn,
            Lsn lastCommitLsn,
            Long txId,
            Instant time,
            boolean snapshot,
            boolean lastSnapshotRecord,
            TransactionContext transactionContext,
            IncrementalSnapshotContext<TableId> incrementalSnapshotContext) {
        this.streamingStoppingLsn = null;
        this.sourceInfo = new SourceInfo(connectorConfig);
        this.lastCompletelyProcessedLsn = lastCompletelyProcessedLsn;
        this.lastCommitLsn = lastCommitLsn;
        this.sourceInfo.update(lsn, time, txId, this.sourceInfo.xmin(), (TableId) null);
        this.sourceInfo.updateLastCommit(lastCommitLsn);
        this.sourceInfoSchema = this.sourceInfo.schema();
        this.lastSnapshotRecord = lastSnapshotRecord;
        if (this.lastSnapshotRecord) {
            this.postSnapshotCompletion();
        } else {
            this.sourceInfo.setSnapshot(snapshot ? SnapshotRecord.TRUE : SnapshotRecord.FALSE);
        }

        this.transactionContext = transactionContext;
        this.incrementalSnapshotContext = incrementalSnapshotContext;
    }

    public Map<String, ?> getOffset() {
        Map<String, Object> result = new HashMap();
        if (this.sourceInfo.timestamp() != null) {
            result.put("ts_usec", Conversions.toEpochMicros(this.sourceInfo.timestamp()));
        }

        if (this.sourceInfo.txId() != null) {
            result.put("txId", this.sourceInfo.txId());
        }

        if (this.sourceInfo.lsn() != null) {
            result.put("lsn", this.sourceInfo.lsn().asLong());
        }

        if (this.sourceInfo.xmin() != null) {
            result.put("xmin", this.sourceInfo.xmin());
        }

        if (this.sourceInfo.isSnapshot()) {
            result.put("snapshot", true);
            result.put("last_snapshot_record", this.lastSnapshotRecord);
        }

        if (this.lastCompletelyProcessedLsn != null) {
            result.put("lsn_proc", this.lastCompletelyProcessedLsn.asLong());
        }

        if (this.lastCommitLsn != null) {
            result.put("lsn_commit", this.lastCommitLsn.asLong());
        }

        return (Map)
                (this.sourceInfo.isSnapshot()
                        ? result
                        : this.incrementalSnapshotContext.store(
                                this.transactionContext.store(result)));
    }

    public Schema getSourceInfoSchema() {
        return this.sourceInfoSchema;
    }

    public Struct getSourceInfo() {
        return this.sourceInfo.struct();
    }

    public boolean isSnapshotRunning() {
        return this.sourceInfo.isSnapshot();
    }

    public void preSnapshotStart() {
        this.sourceInfo.setSnapshot(SnapshotRecord.TRUE);
        this.lastSnapshotRecord = false;
    }

    public void preSnapshotCompletion() {
        this.lastSnapshotRecord = true;
    }

    public void postSnapshotCompletion() {
        this.sourceInfo.setSnapshot(SnapshotRecord.FALSE);
    }

    public void updateWalPosition(
            Lsn lsn,
            Lsn lastCompletelyProcessedLsn,
            Instant commitTime,
            Long txId,
            Long xmin,
            TableId tableId) {
        this.lastCompletelyProcessedLsn = lastCompletelyProcessedLsn;
        this.sourceInfo.update(lsn, commitTime, txId, xmin, tableId);
    }

    public void updateWalPosition(
            Lsn lsn, Lsn lastCompletelyProcessedLsn, Instant commitTime, Long txId, Long xmin) {
        this.updateWalPosition(
                lsn, lastCompletelyProcessedLsn, commitTime, txId, xmin, (TableId) null);
    }

    public void updateCommitPosition(Lsn lsn, Lsn lastCompletelyProcessedLsn) {
        this.lastCompletelyProcessedLsn = lastCompletelyProcessedLsn;
        this.lastCommitLsn = lsn;
        this.sourceInfo.updateLastCommit(lsn);
    }

    boolean hasLastKnownPosition() {
        return this.sourceInfo.lsn() != null;
    }

    boolean hasCompletelyProcessedPosition() {
        return this.lastCompletelyProcessedLsn != null;
    }

    Lsn lsn() {
        return this.sourceInfo.lsn();
    }

    Lsn lastCompletelyProcessedLsn() {
        return this.lastCompletelyProcessedLsn;
    }

    Lsn lastCommitLsn() {
        return this.lastCommitLsn;
    }

    Lsn getStreamingStoppingLsn() {
        return this.streamingStoppingLsn;
    }

    public void setStreamingStoppingLsn(Lsn streamingStoppingLsn) {
        this.streamingStoppingLsn = streamingStoppingLsn;
    }

    Long xmin() {
        return this.sourceInfo.xmin();
    }

    public String toString() {
        return "PostgresOffsetContext [sourceInfoSchema="
                + this.sourceInfoSchema
                + ", sourceInfo="
                + this.sourceInfo
                + ", lastSnapshotRecord="
                + this.lastSnapshotRecord
                + ", lastCompletelyProcessedLsn="
                + this.lastCompletelyProcessedLsn
                + ", lastCommitLsn="
                + this.lastCommitLsn
                + ", streamingStoppingLsn="
                + this.streamingStoppingLsn
                + ", transactionContext="
                + this.transactionContext
                + ", incrementalSnapshotContext="
                + this.incrementalSnapshotContext
                + "]";
    }

    public static PostgresOffsetContext initialContext(
            PostgresConnectorConfig connectorConfig,
            GaussDBConnection jdbcConnection,
            Clock clock) {
        return initialContext(connectorConfig, jdbcConnection, clock, (Lsn) null, (Lsn) null);
    }

    public static PostgresOffsetContext initialContext(
            PostgresConnectorConfig connectorConfig,
            GaussDBConnection jdbcConnection,
            Clock clock,
            Lsn lastCommitLsn,
            Lsn lastCompletelyProcessedLsn) {
        try {
            LOGGER.info("Creating initial offset context");
            Lsn lsn = Lsn.valueOf(jdbcConnection.currentXLogLocation());
            Long txId = jdbcConnection.currentTransactionId();
            LOGGER.info("Read xlogStart at '{}' from transaction '{}'", lsn, txId);
            return new PostgresOffsetContext(
                    connectorConfig,
                    lsn,
                    lastCompletelyProcessedLsn,
                    lastCommitLsn,
                    txId,
                    clock.currentTimeAsInstant(),
                    false,
                    false,
                    new TransactionContext(),
                    new SignalBasedIncrementalSnapshotContext(false));
        } catch (SQLException var7) {
            SQLException e = var7;
            throw new ConnectException("Database processing error", e);
        }
    }

    public OffsetState asOffsetState() {
        return new OffsetState(
                this.sourceInfo.lsn(),
                this.sourceInfo.txId(),
                this.sourceInfo.xmin(),
                this.sourceInfo.timestamp(),
                this.sourceInfo.isSnapshot());
    }

    public void markLastSnapshotRecord() {
        this.sourceInfo.setSnapshot(SnapshotRecord.LAST);
    }

    public void event(DataCollectionId tableId, Instant instant) {
        this.sourceInfo.update(instant, (TableId) tableId);
    }

    public TransactionContext getTransactionContext() {
        return this.transactionContext;
    }

    public void incrementalSnapshotEvents() {
        this.sourceInfo.setSnapshot(SnapshotRecord.INCREMENTAL);
    }

    public IncrementalSnapshotContext<?> getIncrementalSnapshotContext() {
        return this.incrementalSnapshotContext;
    }

    public static class Loader implements OffsetContext.Loader<PostgresOffsetContext> {
        private final PostgresConnectorConfig connectorConfig;

        public Loader(PostgresConnectorConfig connectorConfig) {
            this.connectorConfig = connectorConfig;
        }

        private Long readOptionalLong(Map<String, ?> offset, String key) {
            Object obj = offset.get(key);
            return obj == null ? null : ((Number) obj).longValue();
        }

        public PostgresOffsetContext load(Map<String, ?> offset) {
            Lsn lsn = Lsn.valueOf(this.readOptionalLong(offset, "lsn"));
            Lsn lastCompletelyProcessedLsn = Lsn.valueOf(this.readOptionalLong(offset, "lsn_proc"));
            Lsn lastCommitLsn = Lsn.valueOf(this.readOptionalLong(offset, "lsn_commit"));
            if (lastCommitLsn == null) {
                lastCommitLsn = lastCompletelyProcessedLsn;
            }

            Long txId = this.readOptionalLong(offset, "txId");
            Instant useconds = Conversions.toInstantFromMicros((Long) offset.get("ts_usec"));
            boolean snapshot = Boolean.FALSE;
            boolean lastSnapshotRecord = Boolean.FALSE;
            if (offset.containsKey("snapshot")) {
                snapshot = (Boolean) offset.get("snapshot");
            }
            if (offset.containsKey("last_snapshot_record")) {
                lastSnapshotRecord = (Boolean) offset.get("last_snapshot_record");
            }
            return new PostgresOffsetContext(
                    this.connectorConfig,
                    lsn,
                    lastCompletelyProcessedLsn,
                    lastCommitLsn,
                    txId,
                    useconds,
                    snapshot,
                    lastSnapshotRecord,
                    TransactionContext.load(offset),
                    SignalBasedIncrementalSnapshotContext.load(offset, false));
        }
    }
}
